#!/usr/bin/env python3
# _*_ coding:utf-8 _*_

import argparse

import torch
import torch.nn as nn
import numpy as np


class Options():

    def __init__(self):
        parser = argparse.ArgumentParser()

        parser.add_argument('--no_cuda',
                            action='store_true',
                            default=False,
                            help='Disables cuda training')
        parser.add_argument('--n_gauss',
                            type=int,
                            default=100,
                            help='number of given gauss points')
        parser.add_argument('--n_pde',
                            type=int,
                            default=5000,
                            help='number of interior samples when using LHS for train')
        parser.add_argument('--n_bc',
                            type=int,
                            default=100,
                            help='number of boundary samples when using LHS for train')
        parser.add_argument('--case',
                            type=int,
                            default=1,
                            help='problem case')
        parser.add_argument('--nt',
                            type=int,
                            default=100,
                            help='number of samples in t-direction for train')
        parser.add_argument('--nt_valid',
                            type=int,
                            default=50,
                            help='number of samples in t-direction for valid')
        parser.add_argument('--nx',
                            type=int,
                            default=100,
                            help='number of samples in x-direction for train')
        parser.add_argument('--nx_valid',
                            type=int,
                            default=50,
                            help='number of samples in x-direction for valid')
        parser.add_argument('--nx_plot',
                            type=int,
                            default=100,
                            help='number of samples in x-direction for plot')
        parser.add_argument('--n_valid',
                            type=int,
                            default=200,
                            help='number of interior samples when using LHS for valid')
        parser.add_argument('--n_plot',
                            type=int,
                            default=100,
                            help='number of interior samples when using LHS for valid')
        parser.add_argument('--lr',
                            type=float,
                            default=0.001,
                            help='learning rate')
        parser.add_argument('--epochs_adam',
                            type=int,
                            default=2000,
                            help='epochs for Adam')
        parser.add_argument('--epochs_lbfgs',
                            type=int,
                            default=100,
                            help='epochs for LBFGS')
        parser.add_argument('--tt',
                            type=float,
                            default=0.75,
                            help='test points which t=tt')
        parser.add_argument('--domain',
                            type=list,
                            default=[0, 1, 0, 1],
                            help='[tmin,tmax,xmin,xmax]')
        parser.add_argument('--resume',
                            type=str,
                            default=None,
                            help='put the path to resuming file if needed')
        parser.add_argument('--strategy',
                            type=str,
                            default='LHS',
                            help='Sampling strategy')
        parser.add_argument('--lam',
                            type=float,
                            default=1,
                            help='soft punish cofficient')
        parser.add_argument('--nu',
                            type=float,
                            default=0.001/np.pi,
                            help='viscosity cofficient')
        parser.add_argument('--model',
                            type=str,
                            default='SoftNN',
                            help='model for train (SoftNN, SoftResNet)')
        parser.add_argument('--act_name',
                            type=str,
                            default='tanh',
                            help='activation')
        parser.add_argument('--cuda_id',
                            type=int,
                            default=0,
                            help='device name')
        self.parser = parser

    def parse(self):
        args = self.parser.parse_args()
        args.cuda = not args.no_cuda and torch.cuda.is_available()
        if args.cuda:
            args.device = torch.device(f"cuda:{args.cuda_id}")
        else:
            args.device = torch.device("cpu")

        return args


if __name__ == '__main__':
    args = Options().parse()
    args.epochs_adam = 100
    print(args)
